/**************************************************************************
    filename:   SocketServerDemo.cpp 
    author:     xybei
    created:    2013/08/27
    purpose:    Implement a socket server demo.
                
**************************************************************************/

#define WIN32_LEAN_AND_MEAN

#include <windows.h>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <stdlib.h>
#include <stdio.h>
#include <queue>

using namespace std;

// Need to link with Ws2_32.lib
#pragma comment (lib, "Ws2_32.lib")

#define DEFAULT_PORT "8000"
#define DEFAULT_BUFLEN 128


HANDLE g_hMtxQueue = NULL;
HANDLE g_hSemNumRequests = NULL;

queue<SOCKET> g_queRequestSocket;

DWORD WINAPI AnswerThread(LPVOID lparam);
int GetNumberFromUser(char* szPrompting);

int main(void) 
{
    WSADATA wsaData = { 0 };
    int nResult = 0;

    SOCKET ListenSocket = INVALID_SOCKET;
    SOCKET AcceptSocket = INVALID_SOCKET;

    struct addrinfo* pAddrInfo = NULL;
    struct addrinfo hints = { 0 };

    int nMaxThreadNum = 0;
    HANDLE* phThreadArray = NULL; 

    // Initialize Winsocket
    nResult = WSAStartup(MAKEWORD(2,2), &wsaData);
    if (nResult != 0)
    {
        printf("WSAStartup failed with error: %d\n", nResult);
        return 1;
    }

    ZeroMemory(&hints, sizeof(hints));
    hints.ai_family = AF_INET;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;
    hints.ai_flags = AI_PASSIVE;

    // Resolve the server address and port
    nResult = getaddrinfo(NULL, DEFAULT_PORT, &hints, &pAddrInfo);
    if ( nResult != 0 ) 
    {
        printf("getaddrinfo failed with error: %d\n", nResult);
        WSACleanup();
        return 1;
    }

    // Create a SOCKET for connecting to server
    ListenSocket = socket(pAddrInfo->ai_family, pAddrInfo->ai_socktype, pAddrInfo->ai_protocol);
    if (ListenSocket == INVALID_SOCKET) 
    {
        printf("socket failed with error: %ld\n", WSAGetLastError());
        freeaddrinfo(pAddrInfo);
        WSACleanup();
        return 1;
    }

    // Setup the TCP listening socket
    nResult = bind( ListenSocket, pAddrInfo->ai_addr, (int)pAddrInfo->ai_addrlen);
    if (nResult == SOCKET_ERROR) 
    {
        printf("bind failed with error: %d\n", WSAGetLastError());
        freeaddrinfo(pAddrInfo);
        closesocket(ListenSocket);
        WSACleanup();
        return 1;
    }

    freeaddrinfo(pAddrInfo);

    nResult = listen(ListenSocket, SOMAXCONN);
    if (nResult == SOCKET_ERROR)
    {
        printf("listen failed with error: %d\n", WSAGetLastError());
        closesocket(ListenSocket);
        WSACleanup();
        return 1;
    }

    nMaxThreadNum = GetNumberFromUser("Please input work thread number:\n");

    // Create Mutex and Semaphore.
    g_hMtxQueue = CreateMutex(NULL,
                              FALSE,
                              NULL);
    if (NULL == g_hMtxQueue)
    {
        printf("CreateMutex failed with error:%d \n", GetLastError());
        return 1;
    }
    g_hSemNumRequests = CreateSemaphore(NULL,
                                        0,
                                        nMaxThreadNum,
                                        NULL);
    if (NULL == g_hSemNumRequests)
    {
        printf("CreateSemaphore failed with error:%d \n", GetLastError());
        return 1;
    }

    phThreadArray = new HANDLE[nMaxThreadNum];
    if (NULL == phThreadArray)
    {
        printf("New memory failed!");
        return 1;
    }
    memset(phThreadArray, 0, sizeof(HRESULT)*nMaxThreadNum);

    // Create work threads.
    for (int i = 0; i < nMaxThreadNum; i++)
    {
        phThreadArray[i] = CreateThread(NULL,
                                        NULL,
                                        AnswerThread,
                                        NULL,
                                        0,
                                        NULL);  
        if(phThreadArray[i] == NULL)  
        {  
            printf("CreateThread[%d] failed with error:%d\n", i, GetLastError());  
        }  
        else  
        {  
            printf("CreateThread[%d] OK!\n", i);  
        }
    }

    while(1)
    {
        // Accept a client socket
        AcceptSocket = SOCKET_ERROR;
        AcceptSocket = accept(ListenSocket, NULL, NULL);
        if (AcceptSocket == INVALID_SOCKET)
        {
            printf("accept failed with error: %d\n", WSAGetLastError());
            closesocket(ListenSocket);
            WSACleanup();
            return 1;
        }
        printf("Client Connected.\n");  

        // Add a request to the queue.
        BOOLEAN bAddSuccess = FALSE;
        do 
        {
            WaitForSingleObject(g_hMtxQueue, INFINITE);

            if (ReleaseSemaphore(g_hSemNumRequests, 1, NULL))
            {
                g_queRequestSocket.push(AcceptSocket);
                bAddSuccess = TRUE;
            }

            ReleaseMutex(g_hMtxQueue);

        } while (!bAddSuccess);
    }

    // Cleanup
    closesocket(ListenSocket);

    WSACleanup();

    if (NULL == phThreadArray)
    {
        delete [] phThreadArray;
        phThreadArray = NULL;
    }

    return 0;
}

DWORD WINAPI AnswerThread(LPVOID lparam)  
{  
    int nResult;
    int nSendResult;
    char szRecvBuf[DEFAULT_BUFLEN] = { 0 };  
    SOCKET ClientSocket = 0;

    while(1)
    {
        WaitForSingleObject(g_hMtxQueue, INFINITE);

        if (WAIT_OBJECT_0 == WaitForSingleObject(g_hSemNumRequests, 0))
        {
            if (g_queRequestSocket.empty())
            {
                printf("Error!!! The request socket queue is empty!");
                ReleaseMutex(g_hMtxQueue);
            }
            else
            {
                ClientSocket = g_queRequestSocket.front();
                g_queRequestSocket.pop();

                ReleaseMutex(g_hMtxQueue);

                printf("Current work threadId = %d\n", GetCurrentThreadId());

                // Receive until the peer shuts down the connection
                do 
                {
                    nResult = recv(ClientSocket, szRecvBuf, sizeof(szRecvBuf), 0);
                    if (nResult > 0)
                    {
                         printf("Bytes received: %d, %s\n", nResult, szRecvBuf);

                        // Echo the buffer back to the sender
                        nSendResult = send( ClientSocket, szRecvBuf, nResult, 0 );
                        if (nSendResult == SOCKET_ERROR)
                        {
                            printf("send failed with error: %d\n", WSAGetLastError());
                        }
                        printf("Bytes sent: %d\n", nSendResult);
                    }
                    else if (nResult == 0)
                    {
                        printf("Connection closing...\n");
                    }
                    else
                    {
                        printf("recv failed with error: %d\n", WSAGetLastError());
                    }

                } while (nResult > 0);

                closesocket(ClientSocket);
            }
        }
        else
        {
            ReleaseMutex(g_hMtxQueue);
        }
    }

    return  0;  
} 

int GetNumberFromUser(char* szPrompting)
{
    int nInputNumber = -1;

    printf(szPrompting);
    scanf_s("%d", &nInputNumber);

    return nInputNumber;
}
